import random
import time
import numpy as np
from pathlib import Path
from typing import Tuple, List, Dict, Any

# 导入自定义模块
from src.impl.output import get_output_float16
from src.generator_utils import get_random_seed_tensor
from src.evaluator import has_precision_problem_occurred
from src.classification.classification import get_benchmark


class FrameworkTester:
    """深度学习框架数值精度测试器"""

    # 支持的框架列表
    SUPPORTED_FRAMEWORKS = ['tf', 'torch', 'mnn']

    # 变异幅度配置
    MUTATION_MAGNITUDES = [1e-4, 1e-6, 1e-8]

    def __init__(self, shape: Tuple[int, ...], operator: str,
                 error_threshold: float = 0.0003, mode: str = '1',
                 target_cases_per_framework: int = 100,
                 variable_shape: Tuple[int, ...] = None):
        """
        初始化测试器

        参数:
            shape: 输入张量形状
            operator: 要测试的算子名称
            error_threshold: 精度误差阈值
            mode: 距离计算模式 ('1'或'n')
            target_cases_per_framework: 每个框架目标收集的用例数
            variable_shape: 算子额外变量的形状（如权重/偏置）
        """
        self.shape = shape
        self.operator = operator
        self.error_threshold = error_threshold
        self.mode = mode
        self.target_cases = target_cases_per_framework
        self.variable_shape = variable_shape

        # 为每个框架初始化状态
        self.framework_states = {
            framework: {
                'corpus': [],  # 测试用例语料库
                'max_distance': 0,  # 最大误差距离
                'case_count': 0,  # 已收集用例数
                'data_dir': Path('data') / operator / framework  # 数据保存目录
            }
            for framework in self.SUPPORTED_FRAMEWORKS
        }

        # 创建数据目录
        for state in self.framework_states.values():
            state['data_dir'].mkdir(parents=True, exist_ok=True)

        # 统计信息
        self.total_cross_framework_cases = 0
        self.total_test_cases = 0
        self.start_time = None

    def calculate_distance(self, tensor: np.ndarray, framework: str) -> float:
        """
        计算给定框架输出与基准输出的误差距离

        参数:
            tensor: 输入张量
            framework: 框架名称

        返回:
            计算得到的误差距离
        """
        # 验证框架名称
        if framework not in self.SUPPORTED_FRAMEWORKS:
            raise ValueError(f"不支持的框架: {framework}")

        # 获取基准输出和变量信息
        benchmark, variable = get_benchmark(tensor, self.operator)

        # 获取三个框架的float16输出
        tf_output, torch_output, mnn_output, _ = get_output_float16(
            tensor, self.operator, variable
        )

        # 选择对应框架的输出
        framework_outputs = {
            'tf': tf_output.astype('float16'),
            'torch': torch_output.astype('float16'),
            'mnn': mnn_output.astype('float16')
        }
        output = framework_outputs[framework]

        # 计算绝对误差
        absolute_diff = np.maximum(benchmark - output, output - benchmark)

        # 根据模式计算距离
        if self.mode == '1':
            # 平均绝对误差 (MAE)
            return np.sum(absolute_diff) / np.size(absolute_diff)
        elif self.mode == 'n':
            # 最大绝对误差
            return np.max(absolute_diff)
        else:
            raise ValueError(f"不支持的误差计算模式: {self.mode}")

    def mutate_tensor(self, tensor: np.ndarray) -> np.ndarray:
        """
        对输入张量进行微小扰动（变异）

        参数:
            tensor: 输入张量

        返回:
            添加微小扰动后的张量
        """
        # 随机选择扰动幅度
        magnitude = random.choice(self.MUTATION_MAGNITUDES)

        # 生成随机扰动
        disturbance = random.uniform(0, magnitude)

        return tensor + disturbance

    def get_sample_for_framework(self, framework: str) -> np.ndarray:
        """
        为指定框架获取测试样本（从语料库随机选择或生成新样本）

        参数:
            framework: 框架名称

        返回:
            测试样本张量
        """
        state = self.framework_states[framework]

        # 如果语料库不为空，从中随机选择样本
        if state['corpus']:
            return random.choice(state['corpus'])
        else:
            # 否则生成新的随机样本
            return get_random_seed_tensor(self.shape)

    def test_framework(self, framework: str) -> bool:
        """
        对指定框架执行一次测试

        参数:
            framework: 框架名称

        返回:
            bool: 是否发现了新的精度问题用例
        """
        state = self.framework_states[framework]

        # 如果已经收集足够用例，跳过测试
        if state['case_count'] >= self.target_cases:
            return False

        # 获取测试样本并变异
        sample = self.get_sample_for_framework(framework)
        mutated_sample = self.mutate_tensor(sample)

        # 检测精度问题
        precision_problems = has_precision_problem_occurred(
            mutated_sample, self.operator, 'MRE', self.error_threshold
        )

        # 获取框架对应的检测结果索引
        framework_index = self.SUPPORTED_FRAMEWORKS.index(framework)
        has_problem = precision_problems[framework_index]

        if has_problem:
            # 保存用例文件
            filename = state['data_dir'] / f'{state["case_count"]}.npy'
            np.save(filename, mutated_sample)

            # 更新统计信息
            state['case_count'] += 1
            self.total_cross_framework_cases += 1

            # 计算当前距离
            current_distance = self.calculate_distance(mutated_sample, framework)

            # 如果距离大于当前最大值，添加到语料库
            if current_distance > state['max_distance']:
                state['max_distance'] = current_distance
                state['corpus'].append(mutated_sample)

            return True

        return False

    def print_progress(self):
        """打印测试进度"""
        if self.total_test_cases % 100 == 0:
            print(f'已完成 {self.total_test_cases} 个测试用例')

            # 打印各框架进度
            for framework in self.SUPPORTED_FRAMEWORKS:
                state = self.framework_states[framework]
                print(f'  {framework}: {state["case_count"]}/{self.target_cases}')

    def all_frameworks_completed(self) -> bool:
        """检查是否所有框架都已完成测试"""
        return all(
            state['case_count'] >= self.target_cases
            for state in self.framework_states.values()
        ) or self.total_test_cases > 10000

    def run(self) -> Dict[str, Any]:
        """
        运行测试

        返回:
            包含测试结果的字典
        """
        print(f"开始测试算子: {self.operator}")
        print(f"目标每个框架收集 {self.target_cases} 个用例")
        print("-" * 50)

        self.start_time = time.time()
        self.total_test_cases = 0

        # 主测试循环
        while not self.all_frameworks_completed():
            self.print_progress()

            # 对每个框架执行测试
            for framework in self.SUPPORTED_FRAMEWORKS:
                self.test_framework(framework)

            self.total_test_cases += 1

        # 计算总耗时
        total_time = time.time() - self.start_time

        # 准备结果
        result = {
            'total_time': total_time,
            'total_test_cases': self.total_test_cases,
            'total_cross_framework_cases': self.total_cross_framework_cases,
            'framework_results': {}
        }

        # 添加各框架结果
        for framework in self.SUPPORTED_FRAMEWORKS:
            state = self.framework_states[framework]
            result['framework_results'][framework] = {
                'case_count': state['case_count'],
                'corpus_size': len(state['corpus']),
                'max_distance': state['max_distance']
            }

        return result

    def print_results(self, result: Dict[str, Any]):
        """打印测试结果"""
        print("\n" + "=" * 50)
        print("测试完成!")
        print("=" * 50)

        print(f"总执行时间: {result['total_time']:.2f} 秒")
        print(f"总测试用例数: {result['total_test_cases']}")
        print(f"跨框架精度问题用例总数: {result['total_cross_framework_cases']}")

        print("\n各框架详细结果:")
        for framework, stats in result['framework_results'].items():
            print(f"  {framework.upper()}:")
            print(f"    收集用例数: {stats['case_count']}")
            print(f"    语料库大小: {stats['corpus_size']}")
            print(f"    最大误差距离: {stats['max_distance']:.6f}")


def main():
    """主函数"""
    # 算子配置列表
    OPERATOR_CONFIGS = [
        # 池化操作
        # {'operator': 'avg_pool', 'shape': (4, 3, 14, 14)},
        # {'operator': 'max_pool', 'shape': (4, 3, 14, 14)},

        # 卷积相关
        # {'operator': 'conv2d', 'shape': (4, 3, 14, 14), 'variable_shape': (3, 3, 3, 32)},
        # {'operator': 'bias_add', 'shape': (4, 32, 12, 12), 'variable_shape': (32,)},

        # 全连接层
        # {'operator': 'dense', 'shape': (4, 16)},

        # 归约操作
        {'operator': 'reduce_mean', 'shape': (4, 3, 14, 14)},
        {'operator': 'reduce_max', 'shape': (4, 3, 14, 14)},

        # 激活函数
        {'operator': 'softmax', 'shape': (4, 10)},
        {'operator': 'relu', 'shape': (4, 64)},
        {'operator': 'sigmoid', 'shape': (4, 1)},
        {'operator': 'tanh', 'shape': (4, 64)},

        # 批归一化
        {'operator': 'batch_normalization', 'shape': (4, 3, 14, 14)}
    ]

    # 遍历所有算子配置
    for config in OPERATOR_CONFIGS:
        print(f"\n{'=' * 50}")
        print(f"开始测试算子: {config['operator']}")
        print(f"{'=' * 50}")

        # 创建测试器并运行测试
        tester = FrameworkTester(
            shape=config['shape'],
            operator=config['operator'],
            error_threshold=0.00015,
            mode='n',
            target_cases_per_framework=100,  # 减少测试用例数量以加快测试
            variable_shape=config.get('variable_shape')
        )

        result = tester.run()

        # 打印结果
        tester.print_results(result)


if __name__ == '__main__':
    main()